'''
Created on Sep 28, 2009

@author: mkiyer
'''

from veggie.gene.gene import parse_refseq_genes
from veggie.sample.samplegroup import parse_samplegroups_xml
from optparse import OptionParser
import h5py
import numpy as np
import veggie.coverage.covdb as covdb
import os
import sys

if __name__ == '__main__':
    optionparser = OptionParser("usage: %prog [options]")
    optionparser.add_option("--covdb", dest="covdb",
                            help="path to coverage database")
    optionparser.add_option("-s", "--samples", dest="samples",
                            help="sample XML file to process")
    optionparser.add_option("-o", "--output-dir", dest="output_dir",
                            default=os.path.curdir,
                            help="path to place the output files")
    (options, args) = optionparser.parse_args()


    # open the coverage database file
    h5file = h5py.File(options.covdb, 'r')

    # generate SampleGroupCoverage objects that pair samples with
    # coverage data
    sample_groups = [covdb.SampleGroupCoverage(sgroup, h5file) for sgroup in 
                     list(parse_samplegroups_xml(options.samples))]

    # read the genes
    genes = []    
    # stats
    num_without_exons = 0
    num_without_cds = 0
    num_with_noncoding_utr3_exons = 0
    num_total_genes = 0
    
    for g in parse_refseq_genes():
        num_total_genes += 1
        # ensure gene has exons
        if g.exon_count == 0:
            num_without_exons += 1
            continue
        # ensure gene codes for protein
        if g.cds_start == g.cds_end:
            num_without_cds += 1
            continue
        # get the coding part of the 3'UTR exon
        if g.strand == '+':
            cds_utr3_exon = (g.exons[-1][0], g.cds_end)            
        elif g.strand == '-':
            cds_utr3_exon = (g.cds_start, g.exons[0][1])
        else:
            raise ValueError("gene strand must be either + or -")
        if cds_utr3_exon[0] >= cds_utr3_exon[1]:
            num_with_noncoding_utr3_exons += 1
            continue         
        genes.append(g)

    print 'total genes:', num_total_genes
    print 'genes without exons:', num_without_exons
    print 'genes without coding sequence:', num_without_cds
    print 'genes with >1 3UTR exon', num_with_noncoding_utr3_exons
    print 'total usable genes:', len(genes)
    
    for sgroup in sample_groups:
        for sample_name, sample_dsets in sgroup.iteritems():
            
            outfhd = open(os.path.join(options.output_dir, '%s_3utr_exon_profile.txt' % sample_name), 'w')

            outfhd.write('\t'.join(['#chrom',
                                    'start',
                                    'end',
                                    'name',
                                    'min',
                                    'max',
                                    'mean',
                                    'stdev',
                                    'zscore']))            
            for g in genes:                
                # get the coding part of the 3'UTR exon
                if g.strand == '+':
                    cds_utr3_exon = (g.exons[-1][0], g.cds_end)            
                elif g.strand == '-':
                    cds_utr3_exon = (g.cds_start, g.exons[0][1])
                else:
                    raise ValueError("gene strand must be either + or -")
                
                # get the coverage of 3'UTR coding exon        
                arr = 1000.0 * covdb.calc_avg_cov_array(g.chrom, cds_utr3_exon[0], cds_utr3_exon[1], sample_dsets)
                # compute the min, max, average, and variance of the array
                maxval = np.max(arr)
                minval = np.min(arr)
                meanval = np.mean(arr)
                varval = np.var(arr)
                stdval = np.std(arr)
                zscore = (maxval - meanval) / stdval

                outfhd.write('%s\t%d\t%d\t%s\t%.2f\t%.2f\t%.2f\t%.2f\t%.2f\n' %
                             (g.chrom, 
                              cds_utr3_exon[0], cds_utr3_exon[1], 
                              '%s,%s' % (g.symbol, g.acc), 
                              minval, maxval,
                              meanval, stdval, zscore))
            outfhd.close()

